package com.gmail.woodyc40.common.timing;
import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import java.io.*;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;
/**
* Obtains the instance of instrumentation
*
* @author Pierre C
*/
public class Instrument {
private static Instrumentation instrumentation = null;
public static void agentmain(String string, Instrumentation instrument) {
instrumentation = instrument;
}
public static Instrumentation get() {
if (instrumentation == null) {
try {
attachAgentToJVM(getCurrentPID(), Instrument.class);
} catch (IOException | AttachNotSupportedException | AgentLoadException | AgentInitializationException e) {
e.printStackTrace();
}
}
return instrumentation;
}
//////////////////////
////////////////////// THANKS ICYENE https://github.com/Xyene/ASM-Late-Bind-Agent
//////////////////////
/**
* Gets the current JVM PID.
*
* @return Returns the PID.
*/
public static String getCurrentPID() {
String jvm = ManagementFactory.getRuntimeMXBean().getName();
return jvm.substring(0, jvm.indexOf('@'));
}
public static byte[] getBytesFromStream(InputStream stream) throws IOException {
ByteArrayOutputStream buffer = new ByteArrayOutputStream();
int nRead;
byte[] data = new byte[65536];
while ((nRead = stream.read(data, 0, data.length)) != -1) {
buffer.write(data, 0, nRead);
}
buffer.flush();
return buffer.toByteArray();
}
/**
* Loads an agent into a JVM.
*
* @param agent The main agent class.
* @param resources Array of classes to be included with agent.
* @param pid The ID of the target JVM.
* @throws IOException
* @throws AttachNotSupportedException
* @throws AgentLoadException
* @throws AgentInitializationException
*/
public static void attachAgentToJVM(String pid, Class agent, Class... resources)
throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException {
VirtualMachine vm = VirtualMachine.attach(pid);
vm.loadAgent(generateAgentJar(agent, resources).getAbsolutePath());
vm.detach();
}
/**
* Generates a temporary agent file to be loaded.
*
* @param agent The main agent class.
* @param resources Array of classes to be included with agent.
* @return Returns a temporary jar file with the specified classes included.
* @throws FileNotFoundException
* @throws IOException
*/
public static File generateAgentJar(Class agent, Class... resources) throws IOException {
File jarFile = File.createTempFile("agent", ".jar");
jarFile.deleteOnExit();
Manifest manifest = new Manifest();
Attributes mainAttributes = manifest.getMainAttributes();
// Create manifest stating that agent is allowed to transform classes
mainAttributes.put(Attributes.Name.MANIFEST_VERSION, "1.0");
mainAttributes.put(new Attributes.Name("Agent-Class"), agent.getName());
mainAttributes.put(new Attributes.Name("Can-Retransform-Classes"), "true");
mainAttributes.put(new Attributes.Name("Can-Redefine-Classes"), "true");
JarOutputStream jos = new JarOutputStream(new FileOutputStream(jarFile), manifest);
jos.putNextEntry(new JarEntry(agent.getName().replace('.', '/') + ".class"));
jos.write(getBytesFromStream(agent.getClassLoader().getResourceAsStream(unqualify(agent))));
jos.closeEntry();
for (Class clazz : resources) {
String name = unqualify(clazz);
jos.putNextEntry(new JarEntry(name));
jos.write(getBytesFromStream(clazz.getClassLoader().getResourceAsStream(name)));
jos.closeEntry();
}
jos.close();
return jarFile;
}
private static String unqualify(Class clazz) {
return clazz.getName().replace('.', '/') + ".class";
}
}